#include <ATen/ATen.h>
#include <ATen/Functions.h>
#include <ATen/cuda/CUDAContext.h>

#include <THC/THC.h>
#include <THC/THCDeviceUtils.cuh>

#include <ctime>
#include <vector>
#include <iostream>
#include <curand_kernel.h>


__device__ float _div(const float a, const float b, const bool ceil) {
    if(ceil)
        return THCCeilDiv<float>(a, b);
    else
        return a / b;
}

__device__ float _ratio(float ratio, int n) {
    return ratio + (1.0 / n) - 0.5;
}

__global__ void match_gt_kernel(
    const float *pred_gt_loss,
    const bool *is_foreground,
    const int *seeds,
    const int cols,
    float ratio,
    long *min_gt_inds
) {
    float sum = 0;
    float weights = 0;
    curandState state;
    const int row = blockIdx.x;
    const int inds = row * cols;
    const bool ceil = false;
    int n=0;

    for(int i=0; i<cols; i++) {
        if(is_foreground[inds+i]) {
            sum += pred_gt_loss[inds+i];
            n++;
        }
    }

    if(n == 1) {
        for(int i=0; i<cols; i++) {
            if(is_foreground[inds+i]) {
                min_gt_inds[row] = i;
                return;
            }
        }
    }

    for(int i=0;i<cols;i++) {
        if(is_foreground[inds+i]) {
            weights += _div(sum, pred_gt_loss[inds+i], ceil);
        }
    }
    
    curand_init(seeds[row], weights, 0, &state);
    float randu_f = 0;
    while (!(randu_f > 0)) {
        randu_f = curand_uniform(&state);
    }
    float min_bounded(0.0);
    randu_f *= (weights + min_bounded);
    n = n > 0 ? n : 1;
    // int rand_ui = __float2int_rz(randu_f);
    // printf("seed: %d, weights: %f, sum: %f, rand_f: %f %d \n", seeds[row], weights, sum, randu_f, row);
    
    min_gt_inds[row] = -1;
    for(int i=0; i<cols; i++)
        if(is_foreground[inds+i]) {
            // printf("%f %d \n", _div(sum, pred_gt_loss[inds+i], ceil), row);
            if(randu_f > min_bounded) {
                float w = _div(sum, pred_gt_loss[inds+i], ceil);
                if(ratio > 0.0 && (w / weights) > _ratio(ratio, n)) {
                    min_gt_inds[row] = i;
                    return;
                } else {
                    min_bounded += _div(sum, pred_gt_loss[inds+i], ceil);
                    min_gt_inds[row] = i;
                }
            } else break;
        }
}


at::Tensor rand_match_gt_cuda(const at::Tensor pred_gt_loss, const at::Tensor mask, const float ratio) {
    AT_ASSERTM(pred_gt_loss.type().is_cuda(), "loss must be a CUDA tensor");
    
    int rows = pred_gt_loss.size(0);
    int length = pred_gt_loss.size(1);
    // auto is_foreground = pred_gt_loss != non_label;
    auto min_gt_inds = at::empty({rows}, pred_gt_loss.options().dtype(at::kLong));
    auto seeds = at::randint(0, int(time(NULL)), {rows}, at::CUDA(at::kInt));
    
    float *pred_gt_loss_ = pred_gt_loss.data<float>();
    bool *is_foreground_ = mask.data<bool>();
    int* seeds_ = seeds.data<int>();
    
    dim3 blocks(rows);
    dim3 threads(1);
    match_gt_kernel<<<blocks, threads>>>(
        pred_gt_loss_,
        is_foreground_,
        seeds_,
        length,
        ratio,
        min_gt_inds.data<long>()
    );
    
    return min_gt_inds;
}